#include <iostream>
#include <vector>
#include <fstream>
#include <cstring>
#include <queue>
#include <omp.h>
#include <unordered_set>
#include <random>

#undef NDEBUG

#include <cassert>

namespace KNNOverlap {
    std::vector<std::vector<int>> read_graph(char *path) {
        char directory[256];
        std::sprintf(directory, "%s/graph.graph", path);
        std::ifstream instream;
        instream.open(directory, std::ios::in);
        if (!instream.is_open()) {//判断文件是否打开
            std::cout << "Error opening file" << std::endl;
        }

        int n_item;
        instream >> n_item;
        std::vector<std::vector<int>> final_graph(n_item);
        for (int i = 0; i < n_item; i++) {
            int n_neighbor;
            instream >> n_neighbor;
            final_graph[i].resize(n_neighbor);
            for (int j = 0; j < n_neighbor; j++) {
                instream >> final_graph[i][j];
            }
        }
        instream.close();
        return final_graph;
    }

    std::vector<int>
    get_partition_l(std::vector<std::vector<int>> graph, int n_cluster, float size_variance, int seed) {
        int n_item = graph.size();
        std::vector<int> partition_l(n_item);
        for (int i = 0; i < n_item; i++) {
            partition_l[i] = -1;
        }

        srand(seed);
#pragma omp parallel for
        for (int i = 0; i < n_cluster; i++) {
            int label_max_size = (int) (size_variance * n_item / n_cluster);
            int label_cnt = 0;
            int rand_point = rand() % n_item;
            std::queue<int> bfs_queue;
            bfs_queue.push(rand_point);

            std::unordered_set<int> visited_element;
            visited_element.insert(rand_point);
            label_cnt++;
#pragma omp critical
            {
                partition_l[rand_point] = i;
            };
            while ((int) bfs_queue.size() > 0 && label_cnt < label_max_size) {
                int tmp_point = bfs_queue.front();
                bfs_queue.pop();
                int n_neighbor = graph[tmp_point].size();
                for (int j = 0; j < n_neighbor; j++) {
                    int candidate = graph[tmp_point][j];
                    if (visited_element.find(candidate) != visited_element.end()) {
                        continue;
                    }
#pragma omp critical
                    {
                        if (partition_l[candidate] == -1) {
                            bfs_queue.push(candidate);
                            partition_l[candidate] = i;
                            label_cnt++;
                            visited_element.insert(candidate);
                        }
                    };
                }

            }
        }

        for (int i = 0; i < n_item; i++) {
            if (partition_l[i] != -1) {
                continue;
            }
            std::unordered_set<int> visited_element;
            visited_element.insert(i);

            std::queue<int> bfs_queue;
            bfs_queue.push(i);

            bool has_cluster = false;

            while (!has_cluster) { // the size of bfs_queue can not be 0, bfs should stop in the middle condition
                int tmp_point = bfs_queue.front();
                bfs_queue.pop();
                int n_neighbor = graph[tmp_point].size();
                for (int j = 0; j < n_neighbor; j++) {
                    int candidate = graph[tmp_point][j];
                    if (visited_element.find(candidate) == visited_element.end()) {
                        if (partition_l[candidate] != -1) {
                            has_cluster = true;
                            partition_l[i] = partition_l[candidate];
                        } else {
                            bfs_queue.push(candidate);
                            visited_element.insert(candidate);
                        }
                    }
                }
            }
        }
        return partition_l;
    }

    void test_partition_l(std::vector<int> partition_l, int n_cluster) {
        int n_item = partition_l.size();
        bool is_all_allocate = true;
        for (int i = 0; i < n_item; i++) {
            assert(partition_l[i] != -1);
            if (partition_l[i] == -1) {
                is_all_allocate = false;
                break;
            }
        }
        if (!is_all_allocate) {
            printf("partition_l not all allocate\n");
            exit(-1);
        }
    }

    std::vector<std::vector<float >>
    get_centroids(float *base, std::vector<int> partition_l, int vec_dim, int n_cluster) {
        int n_item = partition_l.size();
        std::vector<std::vector<int>> cluster_idx_l_l(n_cluster);
        for (int i = 0; i < n_item; i++) {
            int cluster_i = partition_l[i];
            cluster_idx_l_l[cluster_i].push_back(i);
        }

        std::vector<std::vector<float>> centroid_l(n_cluster);

#pragma omp parallel for
        for (int i = 0; i < n_cluster; i++) {
            double *centroid_sum = new double[vec_dim];
            for (int j = 0; j < vec_dim; j++) {
                centroid_sum[j] = 0;
            }
            int cluster_idx_size = cluster_idx_l_l[i].size();
            for (int j = 0; j < cluster_idx_size; j++) {
                int idx = cluster_idx_l_l[i][j];
                float *tmp_item = base + idx * vec_dim;
                for (int k = 0; k < vec_dim; k++) {
                    centroid_sum[k] += tmp_item[k];
                }
            }
            std::vector<float> tmp_centroid(vec_dim);
            for (int j = 0; j < vec_dim; j++) {
                tmp_centroid[j] = centroid_sum[j] / cluster_idx_size;
            }
            centroid_l[i] = tmp_centroid;
        }
        return centroid_l;
    }

    void save_centroid(char *centroid_fname, std::vector<std::vector<float >> centroids) {
        int n_cluster = centroids.size();
        int dim = centroids[0].size();

        FILE *fp = NULL;

        fp = fopen(centroid_fname, "w");
        for (int i = 0; i < n_cluster; i++) {
            for (int j = 0; j < dim; j++) {
                fprintf(fp, "%.3f ", centroids[i][j]);
            }
            fprintf(fp, "\n");
        }
        fclose(fp);
    }

}